from collections import deque
from typing import Callable, Deque, List, Optional

import coin2.service.account.account_info_service_pb2 as ais_pb2 
from coin.support.data_replay.reader.interface import (
    RawData,
    IRawDataReader)
from xunkemgmt_client.client.api_client import XunkemgmtClient


class BalanceRealtimeHistoriesGrpcReader(IRawDataReader):
  def __init__(self,
               start_timestamp: int,
               end_timestamp: int,
               grpc_request: ais_pb2.QueryBalanceRealtimeHistoriesRequestProto):
    self._record_cache = self._load_data(
        start_timestamp=start_timestamp,
        end_timestamp=end_timestamp,
        grpc_request=grpc_request)

  def _load_data(self, *, start_timestamp: int, end_timestamp: int, grpc_request
  ) -> Deque[RawData]:
    with XunkemgmtClient() as client:
      rsp = client.query_balance_realtime_histories(grpc_request)
    processed = []
    for history in rsp.histories:
      timestamp = history.balance_timestamp
      if not start_timestamp <= timestamp < end_timestamp:
        continue
      processed.append(RawData(timestamp=timestamp, value=history))
    return deque(sorted(processed, key=lambda x: x.timestamp))

  def peek(self) -> Optional[RawData]:
    value = self._record_cache[0] if len(self._record_cache) > 0 else None
    return value

  def read(self) -> Optional[RawData]:
    value = self._record_cache.popleft() if len(self._record_cache) > 0 else None
    return value


if __name__ == '__main__':
  import datetime
  from coin.base.datetime_util import to_datetime, to_timestamp_int
  start_timestamp = to_timestamp_int(datetime.datetime(2023, 6, 17))
  end_timestamp = to_timestamp_int(datetime.datetime(2023, 6, 17, 1))
  grpc_request = ais_pb2.QueryBalanceRealtimeHistoriesRequestProto(
      market_type='Futures',
      exchange='Binance',
      owner='xunke00_swap',
      start_timestamp=start_timestamp,
      end_timestamp=end_timestamp)
  reader = BalanceRealtimeHistoriesGrpcReader(
      start_timestamp, end_timestamp, grpc_request)
  prev_ts = 0
  while True:
    record = reader.read()
    if record is None:
      break
    curr_ts = record.timestamp
    assert curr_ts >= prev_ts, (prev_ts, curr_ts)
    prev_ts = curr_ts
    print(record.value)
